package main

import (
	"flag"
	"fmt"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"log"
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"
	"tool"
)

var (
	portProbe = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "port_probe_status",
			Help: "Status of remote ports",
			//
			// ConstLabels: prometheus.Labels{"host_ip": getHostIP("HOSTIP")},
		},
		// 通过label实现区分.
		// 这个其实就是我们需要的标签信息.我们在进行赋值的时候,也需要满足这个顺序.
		[]string{"host_ip", "target"})
)

func init() {
	prometheus.MustRegister(portProbe)
}

func main() {
	// 添加 flag ,实现 -c 指定配置文件.
	configfile := flag.String("c", "config.yml", "config file")
	// 内部抓取时间间隔,不受Prometheus的抓取时间影响.
	intervel := flag.Int64("i", 30, "Grab interval to collect")
	// 添加动态端口
	port := flag.String("p", "8889", "Server Port")
	// 超时探测时间
	probetimeout := flag.Int64("t", 2, "Probe TimeOut")

	flag.Parse()

	log.Println("Config file:", *configfile)
	log.Printf("period  Seconds : %d 秒 \n", *intervel)
	log.Printf("timeout Seconds : %d 秒 \n", *probetimeout)
	log.Printf("Server Port: %s \n", *port)

	http.Handle("/metrics", promhttp.Handler())
	intervalTime := time.Duration(*intervel) * time.Second
	// 开启http服务
	go func() {
		log.Println("Starting....")
		err := http.ListenAndServe(":"+*port, nil)
		if err != nil {
			log.Println("Failed to start HTTP server:", err)
		}
	}()

	// 读取配置文件.
	probe := tool.ReadYamlFile(*configfile)
	// 1.创建需要探测的List
	ListProbe := getList(probe)
	// 2. 开启线程,进行端口探测.

	var wg sync.WaitGroup
	// 获取本机的物理ip.注意如果是容器,需要想办法设置这个环境变量.
	hostIp := getHostIP("HOSTIP")

	for {
		// 逐个端口探测.
		for _, target := range ListProbe {
			wg.Add(1)
			go func(target string) {
				defer wg.Done()
				portProbe.WithLabelValues(hostIp, target).Set(ProbePort("tcp", target, *probetimeout))
			}(target)
		}
		wg.Wait()
		// 主线程休眠时间.
		time.Sleep(intervalTime)
	}

}

func getList(probe tool.Probe) []string {
	var target string
	list := []string{}

	for _, schema := range probe.SchemaList {
		for _, host := range schema.Hosts {
			for _, port := range host.Ports {
				target = host.Host + ":" + strconv.Itoa(port)
				list = append(list, target)
			}
		}
	}
	return list
}

// 获取宿主机IP
func getHostIP(key string) string {
	return os.Getenv(key)
}

// ProbePort 探测端口是否通断.
// 端口通,返回true, 不通,返回false
func ProbePort(protocal string, addr string, timeout int64) float64 {
	//duration := time.Duration(timeout)
	startTime := time.Now()
	// 探测端口,默认超时时间是 2秒. 如果有必要,可以调整.
	conn, err := net.DialTimeout(protocal, addr, time.Duration(timeout*int64(time.Second)))
	if err != nil {
		log.Printf("%s is unreachable, after %s probe timeout \n", addr, time.Since(startTime))
		return 0
	} else {
		defer conn.Close()
		// 端口正常,不输出日志
		// log.Printf("%s is reachable, probe spend time %s \n", addr, time.Since(startTime))
		return 1
	}
}

func probePort(targetHost string, targetPort int) time.Duration {
	startTime := time.Now()

	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", targetHost, targetPort), 3*time.Second)
	if err != nil {
		log.Printf("Port %d on %s is unreachable\n", targetPort, targetHost)
	} else {
		defer conn.Close()
		log.Printf("Port %d on %s is reachable\n", targetPort, targetHost)
	}

	return time.Since(startTime)
}
